# 动态规划50题 https://www.bilibili.com/video/BV1aa411f7uT
# 14/50 统计全为1的正方形子矩阵
# leetcode第1277题（同221题）: https://leetcode.cn/problems/count-square-submatrices-with-all-ones/
# Date: 2024/11/4
from leetcode import test
import json


def countSquares(matrix: list[list[int]]) -> int:
    """我的方法，使用一个简单的依次判断所有子方阵的方法，以矩阵的右下角矩阵作为锚点，判断每个元素作为矩阵的右下角元素。
    dp矩阵存放了遍历到第i行j列元素的所包含的全1矩阵的个数。
    本质上这个方法并没有利用动态规划而是穷举法"""

    def element_exists(matrix, target):
        return any(target in row for row in matrix)

    m, n = len(matrix), len(matrix[0])  # 行数, 列数
    dp = [[0] * (n + 1) for _ in range(m)]
    dp[0][0] = 0
    for i in range(m):  # 逐行遍历, 每行的第一列存储上一行最后一列的结果
        for j in range(1, n + 1):
            dp[i][j] = dp[i][j - 1]
            if matrix[i][j - 1] == 1:  # 如果当前遍历的元素为1，则再进行判断以减少时间复杂度
                for k in range(min(i + 1, j), 0, -1):
                    subm = [row[j - k:j] for row in matrix[i - k + 1:i + 1]]
                    if not element_exists(subm, 0):
                        dp[i][j] += 1

        if i + 1 < m:
            dp[i + 1][0] = dp[i][n]  # 每行遍历结束之后, 将最后一个值传递给下一行的第一个值

    return dp[m - 1][n]


def countSquares_dp(matrix: list[list[int]]) -> int:
    """使用二维数组来作为存储变量的动态规划算法。
    这里dp矩阵存放的是以第i行j列元素为锚点，把其当做矩阵的右下角，包含全1方阵的个数.
    如下面的矩阵, 其dp如右侧矩阵:
    [[0, 1, 1, 1],     [[0, 1, 1, 1],
     [1, 1, 1, 1],  ->  [1, 1, 2, 2],
     [0, 1, 1, 1]]      [0, 1, 2, 3]]
    在matrix中对应元素为1, 则判断每个dp元素的值只需要考虑其左、上以及左上元素的最小值再加1即可:
    dp[i][j] = min(dp[i-1][j-1], dp[i-1][j], dp[i][j-1]) + 1
    在matrix中对应元素为0, 则无需判断, 对应的dp元素即为0
    """
    m, n = len(matrix), len(matrix[0])  # 行数, 列数
    dp = [[0] * n for _ in range(m)]
    res = 0
    for i in range(m):
        for j in range(n):
            if matrix[i][j] == 0:
                dp[i][j] = 0
            else:
                dp[i][j] = min(dp[i - 1][j - 1], dp[i][j - 1], dp[i - 1][j]) + 1
            res += dp[i][j]

    return res


def countSquares_dp_3D(matrix: list[list[int]]) -> int:
    """使用三维数组来作为存储变量的动态规划算法。"""
    m, n = len(matrix), len(matrix[0])  # 行数, 列数
    l = max(m, n)
    dp = [[[0] * n for _ in range(m)] for _ in range(l)]


if __name__ == '__main__':
    matrix = [
        [0, 1, 1, 1],
        [1, 1, 1, 1],
        [0, 1, 1, 1]
    ]  # 15

    matrix0 = [[0, 1],
               [1, 0]]  # 2

    matrix1 = [
        [1, 0, 1],
        [1, 1, 0],
        [1, 1, 0]
    ]  # 7

    matrix2 = [
        [0, 0, 1],
        [1, 1, 0],
        [0, 1, 1]
    ]  # 5

    matrix3 = [
        [1, 1, 1],
        [1, 1, 1],
        [1, 1, 1]
    ]  # 14

    matrix4 = [
        [1, 1, 1, 1, 0],  # 23
        [1, 1, 1, 1, 0],  # 14
        [1, 1, 1, 1, 1],  # 7
        [1, 1, 1, 1, 1],  # 2
        [1, 1, 1, 1, 1],
    ]  # 46

    with open("../data/dp14_lc1277_lc221.data", "r") as file:
        lines = file.read().strip()  # 去除首位的空白
    matrix5 = json.loads(lines)

    inp = [{"matrix": matrix}, {"matrix": matrix0}, {"matrix": matrix1}, {"matrix": matrix2}, {"matrix": matrix3},
           {"matrix": matrix4},
           {"matrix": matrix5}, ]
    out = [15, 2, 7, 5, 14, 46, 22859]

    test.test_function(countSquares, inp, out)
    test.test_function(countSquares_dp, inp, out)
